
#############################################################################
# "metasploit" action for driving Metasploit module execution
#############################################################################

# Standard library imports
import sys
import msgpack
import random
import requests
from time import sleep

# Local imports
from database_base import sqlConfig
from storyboard import Storyboard

# Local constants
MODULE_KEY = "module" # Same as defined in Storyboard
USE_KEY = "use"
KEYWORD_KEY = "keyword"
COMMAND_KEY = "command"

COMMENT_SUCCESS = "Metasploit module succeeded"
COMMENT_FAILURE = "Metasploit module failed"
COMMENT_ERROR = "Metasploit module encountered error"

# TODO: Add constants for Metasploit internal strings


class metasploit():

        # Constructor
	def __init__(self):

                # Get Metasploit configuration object
		msf_config = sqlConfig().getMsf()

                # Get Metasploit configuration parameters to use with the RPC API server
		server_ip = msf_config[Storyboard.MSF_SERVER]
		server_port = msf_config[Storyboard.MSF_PORT]
		self.user = msf_config[Storyboard.MSF_USER]
		self.password = msf_config[Storyboard.MSF_PASSWORD]

		# Build the URL for accessing the Metasploit server
		self.url = "http://" + server_ip + ":" + str(server_port) + "/api/"
		# HTTP header
		self.header = {"Content-type":"binary/message-pack"}

		# Token and console id generated by Metasploit
		self.token = ""
		self.console_id = ""

	# Communication with Metasploit
	def api(self, option):
		payload = msgpack.packb(option)

                try:
		        request = requests.post(self.url, headers=self.header, data=payload)
                except Exception as exception:
                        print("[-] metasploit: ERROR: Could not connect with Metasploit: {0}".format(str(exception)))
                        sys.exit(1)

		if request.status_code != 200:
			print("[-] metasploit: ERROR: Could not connect with Metasploit: Status code: {0})".format(request.status_code))
			sys.exit(1)

                # Decode and return the response as text
		return msgpack.unpackb(request.content)

	# Login into Metasploit using the auth.login API
	def msf_login(self):
		option = ["auth.login", self.user, self.password]
		response = self.api(option)

		if "result" in response:
			if response["result"] == "success":
                                # Store the authentication token and return success
				self.token = response["token"]
				return

		print("[-] metasploit: ERROR: Failed to retrieve the authentication token")
		sys.exit()

	# Create an msfconsole instance
	def create_console(self):
		option = ["console.create",self.token]

		while True:
			response = self.api(option)

			if "busy" in response:
				if not response["busy"]:
					break

		# Store the console id
		self.console_id = response["id"]

	# Read from the msfconsole instance
	def read_console(self):
		option = ["console.read", self.token, self.console_id]

		while True:
			response = self.api(option)

			# When 'busy' is present, the command has not completed yet
			if "busy" in response:
				if not response["busy"]:
			                # When 'data' is empty, buffer was not retrieved yet
					if response["data"] != "":
						break

		# Return the data as text
		return response["data"]

	# Execute a command via msfconsole
	def exec_console(self, command):
		# Send the command to msfconsole
		# response = {"wrote":8}
		option = ["console.write", self.token, self.console_id, command + "\n"]
		response = self.api(option)

	# Get list of exploit modules (don't use them)
	def get_exploit_module(self):
		option = ["module.exploits", self.token]

		response = self.api(option)

		# Return the exploit module list
		return response["modules"]

	# Look for modules using the 'search' command with parameter 'keyword'
	def search_module(self, keyword):
		# Execute the 'search' command
		command = "search " + keyword
		self.exec_console(command)

		# Get the command output
		result = self.read_console()

		# Parse the command output
		flag = False
		modules = []
		for line in result.splitlines():
			if line == "":
				continue

			if flag:
				check = line.split()[3]
				if check == "No":
					continue

				modules.append(line.split()[0])

			if "-----" in line:
				flag = True

		# Return the list of matching module
		return modules

	# Get module options
	def get_module_option(self, module):
		module_type = module.split("/")[0]

		option = ["module.options", self.token, module_type, module]
		response = self.api(option)

		# Return options as dictionary
		return response

	# Execute module
	def execute_module(self, module, options):
		module_type = module.split("/")[0]

		option = ["module.execute", self.token, module_type, module, options]
		response = self.api(option)

		# Return result as text
		return response

	# Get jobs
	def get_job(self):
		option = ["job.list", self.token]
		response = self.api(option)

		# Return jobs as dictionary
		return response

	# Get sessions
	def get_session(self):
		option = ["session.list", self.token]
		response = self.api(option)

		# Return sessions as dictionary
		return response

	# Stop session
	def stop_session(self, session_id):
		option = ["session.stop", self.token, session_id]
		response = self.api(option)

		# Return result as text
		return response["result"]

	# Read from session shell
	def session_shell_read(self, session_id):
		option = ["session.shell_read", self.token, session_id]

		for i in range(5):
			response = self.api(option)["data"]

			if response != "":
				break

			sleep(1)

		# Return result as text
		return response

	# Write to session shell
	def session_shell_write(self, session_id, command):
		option = ["session.shell_write", self.token, session_id, command + "\n"]
		response = self.api(option)

		# Return result as dictionary
		return response

	# Execute command via session shell
	def session_shell(self,session_id,command):
		self.session_shell_write(session_id, command)
		return self.session_shell_read(session_id)

	# Prepare options
	def set_option(self, action, module, address):
		# Get module options
		options = self.get_module_option(module)

		# Set standard options
		if "RHOSTS" in options:
			NewOptions = {"RHOSTS":address}
		elif "RHOST" in options:
			NewOptions = {"RHOST":address}

		# Set custom options
		for option in action:

                        # Ignore our own parameters

			if option == MODULE_KEY or option == USE_KEY or option == KEYWORD_KEY or option == COMMAND_KEY:
				continue

                        # Add those options that are present in the valid option list for this module
			if option.upper() in options:
				NewOptions[option.upper()] = action[option]

                # Check whether all required options were provided
		for name, config in options.items():
			if config["required"] == True and "default" not in config:
				if name not in NewOptions:
					print("[-] metasploit: ERROR: Required option for module '{0}' not provided: '{1}'".format(module, name))
					sys.exit()

		return NewOptions

        # Check action parameters
	def check(self,action):

                if USE_KEY in action and KEYWORD_KEY in action:
                        print("[-] metasploit: ERROR: Parameters '{0}' and '{1}' are not compatible".format(USE_KEY, KEYWORD_KEY))
                        return False

                if USE_KEY not in action and KEYWORD_KEY not in action:
                        print("[-] metasploit: ERROR: Either parameter '{0}' or '{1}' must be defined".format(USE_KEY, KEYWORD_KEY))
                        return False

                # Command is not mandatory...
                # if COMMAND_KEY not in action:
                #        print("[-] metasploit: ERROR: Parameter '{0}' is mandatory".format(COMMAND_KEY))
                #	 return False

		return True

        # Execute action
	def action(self, team_name, address, action, data):

		# Get the authentication token
		login = self.msf_login()

		# Create a console and get its id
		self.create_console()
		# Read the ASCII art banner from console
		art = self.read_console()

		# Get the name of the module to use if provided
		if USE_KEY in action:
			module = action[USE_KEY]

		# Otherwise use the keyword to search for matching modules
		elif KEYWORD_KEY in action:
			modules = self.search_module(action[KEYWORD_KEY])

			# If not module was found display an error
			if len(modules) == 0:
				print("[-] metasploit: ERROR: No module name matches keyword '{0}'".format(action[KEYWORD_KEY]))
				sys.exit(-1)

                        # If multiple modules match the keyword, display the list and exit
			elif len(modules) != 1:
				print("[-] metasploit: WARNING: Multiple module names match the keyword '{0}'".format(action[KEYWORD_KEY]))

				print("    --- MATCHING MODULES ---")
				for module in modules:
					print("    {0}".format(module))
				sys.exit()

                        # If only one module matches the keyword, we select it
			module = modules[0]

		NewOptions = self.set_option(action, module, address)

		# Execute module
		attack = self.execute_module(module, NewOptions)

		if "error" in attack:
			for error in attack["error_backtrace"]:
				print(error)

			return False, COMMENT_ERROR, {}

		job = attack["job_id"]
		uuid = attack["uuid"]

		# Wait until all Metasploit jobs finish execution
		while True:
			jobs = self.get_job()
			if str(job) not in jobs.keys():
				break

		sessions = self.get_session()
		for session,info in sessions.items():
			if info["exploit_uuid"] == uuid:
				break
		else:
			return False, COMMENT_FAILURE, {}

                result = {}
		if COMMAND_KEY in action:
			if isinstance(action[COMMAND_KEY], list):
				for command in action[COMMAND_KEY]:
					result = self.session_shell(session, command)
			else:
				result = self.session_shell(session, action[COMMAND_KEY])

		remove = self.stop_session(session)

		return True, COMMENT_SUCCESS, result

	# Destructor
	def __del__(self):
		if self.token != "":
			# Destroy the console
			option = ["console.destroy", self.token, self.console_id]
			res = self.api(option)

			# Remove the authentication token
			option = ["auth.token_remove",self.token,self.token]
			res = self.api(option)


# Use code below to test the module independently
#msf = metasploit()
#option = {"use":"auxiliary/scanner/ssh/ssh_login","USERNAME":"root","PASSWORD":"root"}
#result,comment = msf.action("team-name","192.168.56.102",option)
#del msf
#print comment
